#include <ompl/contrib/hplanner/collision_estimator.h>

namespace ompl
{
namespace base
{

float CollisionStatusEstimator::computeCollisionProb(const State* s, bool* sufficient) const
{
  if(oracle_->getUsePositiveState())
  {
    QueryResultType collision_result, collision_free_result;

    oracle_->pointKNNQuery(s, true, collision_result);
    oracle_->pointKNNQuery(s, false, collision_free_result);

    unsigned int n_collision_knn = collision_result.size();
    unsigned int n_collision_free_knn = collision_free_result.size();

    if(n_collision_knn == 0 && n_collision_free_knn == 0)
    {
      if(sufficient) *sufficient = false;
      return 0.5;
    }

    if(sufficient) *sufficient = true;

    float collision_weight = 0;
    float collision_free_weight = 0;

    for(unsigned int i = 0; i < n_collision_knn; ++i)
    {
      collision_weight += exp(-obstacle_density_ * sqrt(collision_result[i].second));
    }

    for(unsigned int i = 0; i < n_collision_free_knn; ++i)
    {
      collision_free_weight += exp(-obstacle_density_ * sqrt(collision_free_result[i].second));
    }

    return collision_weight / (collision_free_weight + collision_weight);
  }
  else
  {
    QueryResultType collision_result;

    oracle_->pointKNNQuery(s, true, collision_result);

    unsigned int n_collision_knn = collision_result.size();

    if(n_collision_knn == 0)
    {
      if(sufficient) *sufficient = false;
      return 0.5;
    }

    if(sufficient) *sufficient = true;
    float collision_weight = 0;

    for(unsigned int i = 0; i < n_collision_knn; ++i)
    {
      collision_weight += exp(-obstacle_density_ * sqrt(collision_result[i].second));
    }

    return collision_weight;
  }
}


float CollisionStatusEstimator::computeCollisionProbSample(const State* s1, const State* s2, bool* sufficient) const
{
  StateSpace* stateSpace = oracle_->getSpaceInformation()->getStateSpace().get();

  float prob = 0;
  int nd = stateSpace->validSegmentCount(s1, s2);

  if(nd > 1)
  {
    base::State *test = oracle_->getSpaceInformation()->allocState();

    for(int j = 1; j < nd; ++j)
    {
      stateSpace->interpolate(s1, s2, (double)j / (double)nd, test);
      float tmp_prob = computeCollisionProb(test, sufficient);
      if(tmp_prob > prob) prob = tmp_prob;
    }

    oracle_->getSpaceInformation()->freeState(test);
  }

  float tmp_prob = computeCollisionProb(s2, sufficient);
  if(tmp_prob > prob) prob = tmp_prob;

  return tmp_prob;
}

float CollisionStatusEstimator::computeCollisionProbSample(const State* s1, const State* s2, std::vector<float>& probs, bool* sufficient) const
{
  StateSpace* stateSpace = oracle_->getSpaceInformation()->getStateSpace().get();

  float prob = 0;
  int nd = stateSpace->validSegmentCount(s1, s2);

  if(nd > 1)
  {
    base::State *test = oracle_->getSpaceInformation()->allocState();

    for(int j = 1; j < nd; ++j)
    {
      stateSpace->interpolate(s1, s2, (double)j / (double)nd, test);
      float tmp_prob = computeCollisionProb(test, sufficient);
      probs.push_back(tmp_prob);
      if(tmp_prob > prob) prob = tmp_prob;
    }

    oracle_->getSpaceInformation()->freeState(test);
  }

  float tmp_prob = computeCollisionProb(s2, sufficient);
  probs.push_back(tmp_prob);
  if(tmp_prob > prob) prob = tmp_prob;

  return tmp_prob;
}


float CollisionStatusEstimator::computeCollisionProbSampleLine(const State* s1, const State* s2, bool* sufficient) const
{
  QueryResultType collision_result, collision_free_result;

  oracle_->lineKNNQuery(s1, s2, true, collision_result);

  if(oracle_->getUsePositiveState())
    oracle_->lineKNNQuery(s1, s2, false, collision_free_result);

  unsigned int n_collision_knn = collision_result.size();
  unsigned int n_collision_free_knn = collision_free_result.size();

  if(n_collision_knn == 0 && n_collision_free_knn == 0)
  {
    if(sufficient) *sufficient = false;
    return 0.5;
  }

  if(sufficient) *sufficient = true;

  StateSpace* stateSpace = oracle_->getSpaceInformation()->getStateSpace().get();

  float prob = 0;
  int nd = stateSpace->validSegmentCount(s1, s2);

  if(nd > 1)
  {
    base::State *test = oracle_->getSpaceInformation()->allocState();

    for(int j = 1; j < nd; ++j)
    {
      stateSpace->interpolate(s1, s2, (double)j / (double)nd, test);
      float tmp_prob = computeCollisionProbSampleLine_(test, collision_result, collision_free_result);
      if(tmp_prob > prob) prob = tmp_prob;
    }

    oracle_->getSpaceInformation()->freeState(test);
  }

  float tmp_prob = computeCollisionProbSampleLine_(s2, collision_result, collision_free_result);
  if(tmp_prob > prob) prob = tmp_prob;

  return tmp_prob;
}


float CollisionStatusEstimator::computeCollisionProbSampleLine(const State* s1, const State* s2, std::vector<float>& probs, bool* sufficient) const
{
  QueryResultType collision_result, collision_free_result;

  oracle_->lineKNNQuery(s1, s2, true, collision_result);

  if(oracle_->getUsePositiveState())
    oracle_->lineKNNQuery(s1, s2, false, collision_free_result);

  unsigned int n_collision_knn = collision_result.size();
  unsigned int n_collision_free_knn = collision_free_result.size();

  if(n_collision_knn == 0 && n_collision_free_knn == 0)
  {
    if(sufficient) *sufficient = false;
    return 0.5;
  }

  if(sufficient) *sufficient = true;

  StateSpace* stateSpace = oracle_->getSpaceInformation()->getStateSpace().get();

  float prob = 0;
  int nd = stateSpace->validSegmentCount(s1, s2);

  if(nd > 1)
  {
    base::State *test = oracle_->getSpaceInformation()->allocState();

    for(int j = 1; j < nd; ++j)
    {
      stateSpace->interpolate(s1, s2, (double)j / (double)nd, test);
      float tmp_prob = computeCollisionProbSampleLine_(test, collision_result, collision_free_result);
      probs.push_back(tmp_prob);
      if(tmp_prob > prob) prob = tmp_prob;
    }

    oracle_->getSpaceInformation()->freeState(test);
  }

  float tmp_prob = computeCollisionProbSampleLine_(s2, collision_result, collision_free_result);
  probs.push_back(tmp_prob);
  if(tmp_prob > prob) prob = tmp_prob;

  return tmp_prob;
}


std::pair<float, float> CollisionStatusEstimator::computeCollisionSVM(const State* s1, const State* s2, bool* sufficient) const
{
  if(!oracle_->getUsePositiveState())
    return std::make_pair(1.0, 1.0);

  QueryResultType collision_result, collision_free_result;

  unsigned int dim = oracle_->getStateToVectorTransformer().dim();
  float* line_direction = new float[dim];
  float* line_origin = new float[dim];
  float line_len;

  oracle_->lineKNNQuery(s1, s2, true, collision_result, line_origin, line_direction, &line_len);
  oracle_->lineKNNQuery(s1, s2, false, collision_free_result);

  unsigned int n_collision_knn = collision_result.size();
  unsigned int n_collision_free_knn = collision_free_result.size();

  if(n_collision_knn == 0 && n_collision_free_knn == 0)
  {
    if(sufficient) *sufficient = false;
    return std::make_pair(1.0, 1.0);
  }

  if(sufficient) *sufficient = true;

  classifier_svm_->reset();

  for(unsigned int i = 0; i < n_collision_knn; ++i)
  {
    float dist = collision_result[i].second;
    float* item = oracle_->getNegativeRawState(collision_result[i].first);
    classifier_svm_->addItem(item, false);
  }

  for(unsigned int i = 0; i < n_collision_free_knn; ++i)
  {
    float dist = collision_free_result[i].second;
    float* item = oracle_->getPositiveRawState(collision_free_result[i].second);
    classifier_svm_->addItem(item, true);
  }

  classifier_svm_->train();

  float wa = 0;
  float av = 0;
  for(unsigned int i = 0; i < dim; ++i)
  {
    wa += classifier_svm_->w_[i] * line_origin[i];
    av += line_origin[i] * line_direction[i];
  }

  float t = - (wa + classifier_svm_->b_) / av;

  delete [] line_direction;
  delete [] line_origin;

  if(t > 0 && t < line_len)
    return std::make_pair(t / line_len, t);
  else
    return std::make_pair(1.0, 1.0);
}

float CollisionStatusEstimator::computeCollisionProbGen(const State* s1, const State* s2, bool* sufficient) const
{
  QueryResultType collision_result, collision_free_result;

  oracle_->lineKNNQuery(s1, s2, true, collision_result);

  if(oracle_->getUsePositiveState())
    oracle_->lineKNNQuery(s1, s2, false, collision_free_result);

  unsigned int n_collision_knn = collision_result.size();
  unsigned int n_collision_free_knn = collision_free_result.size();

  if(n_collision_knn == 0 && n_collision_free_knn == 0)
  {
    if(sufficient) *sufficient = false;
    return 0.5;
  }

  if(sufficient) *sufficient = true;

  float collision_weight = 0;

  for(unsigned int i = 0; i < n_collision_knn; ++i)
  {
    float dist = collision_result[i].second;
    float tmp_prob = exp(-obstacle_density_ * dist);
    collision_weight += tmp_prob;
  }

  if(!oracle_->getUsePositiveState())
    return collision_weight;

  float collision_free_weight = 0;

  for(unsigned int i = 0; i < n_collision_free_knn; ++i)
  {
    float dist = collision_free_result[i].second;
    float tmp_prob = exp(-obstacle_density_ * dist);
    collision_free_weight += tmp_prob;
  }

  return collision_weight / (collision_weight + collision_free_weight);
}


float CollisionStatusEstimator::computeClosestInCollisionSampleDist(const State* s1, const State* s2) const
{
  QueryResultType collision_result;

  oracle_->lineKNNQuery(s1, s2, true, collision_result);

  if(collision_result.size() > 0)
  {
    return collision_result[0].second;
  }
  else
    return std::numeric_limits<float>::max();
}


float CollisionStatusEstimator::computeCollisionProbSampleLine_(const State* s, const QueryResultType& collision_result, const QueryResultType& collision_free_result) const
{
  unsigned int n_collision_knn = collision_result.size();
  unsigned int n_collision_free_knn = collision_free_result.size();

  unsigned int dim = oracle_->getStateToVectorTransformer().dim();

  if(n_collision_knn == 0 && n_collision_free_knn == 0)
    return 0.5;

  float* sq = oracle_->getStateToVectorTransformer().transform(s);

  float collision_weight = 0;

  for(unsigned int i = 0; i < n_collision_knn; ++i)
  {
    float* d = oracle_->getNegativeRawState(collision_result[i].first);
    float dist = 0;
    for(unsigned int i = 0; i < dim; ++i)
    {
      float tmp = d[i] - sq[i];
      dist += tmp * tmp;
    }
    float tmp_prob = exp(-obstacle_density_ * dist);
    collision_weight += tmp_prob;
  }

  if(!oracle_->getUsePositiveState())
    return collision_weight;

  float collision_free_weight = 0;

  for(unsigned int i = 0; i < n_collision_free_knn; ++i)
  {
    float* d = oracle_->getPositiveRawState(collision_free_result[i].first);
    float dist = 0;
    for(unsigned int i = 0; i < dim; ++i)
    {
      float tmp = d[i] - sq[i];
      dist += tmp * tmp;
    }

    float tmp_prob = exp(-obstacle_density_ * dist);
    collision_free_weight += tmp_prob;
  }

  return collision_weight / (collision_weight + collision_free_weight);
}

}
}
